# _*_coding:utf-8 _*_

from typing import Union
from fastapi import FastAPI
from fastapi.middleware.cors import CORSMiddleware
import os
import traceback
from datetime import datetime
import random
import time
import logging
from logging.handlers import RotatingFileHandler
import coloredlogs
from pydantic import BaseModel, Field

from FlagEmbedding import FlagModel


'''全局参数'''

apikey = 'bge-zh-5f3a8c9b-7e2d-4a1f-9c0d-3b6e8f7a2c1d'
mod_name = '../embd/bge-large-zh-v1.5'


'''定义日志程序'''


'''日志文件名函数'''

def logfilename():
    try:
        nowtime = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime(time.time()))  # 获取当前时间
        pid = os.getpid()  # 获取进程ID
        random_number = round(random.uniform(100, 999))  # 生成随机数
        return f'{nowtime}_{pid}_{random_number}'
    except Exception as e:
        print("处理日志文件名错误:")
        print(e)
        print(traceback.format_exc())
        return ''


'''配置日志文件名、大小、数量、错误日志多文件、颜色'''

try:
    # 创建处理器并设置格式
    formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')

    # 普通日志处理器（INFO及以上，自动轮转）
    file_handler = RotatingFileHandler(f'log/zyzntai_{logfilename()}.log', maxBytes=10*1024*1024,
                                       backupCount=5, encoding='utf-8')
    file_handler.setLevel(logging.INFO)
    file_handler.setFormatter(formatter)

    # 错误日志处理器（ERROR及以上，单独文件）
    error_file_handler = RotatingFileHandler(
        f'log/znt_error{logfilename()}.log', maxBytes=1024*1024, backupCount=5, encoding='utf-8'
    )
    error_file_handler.setLevel(logging.ERROR)
    error_file_handler.setFormatter(formatter)

    # 控制台处理器
    console_handler = logging.StreamHandler()
    console_handler.setLevel(logging.INFO)
    console_handler.setFormatter(formatter)

    # 一次性配置全局日志
    logging.basicConfig(
        level=logging.INFO,
        handlers=[file_handler, error_file_handler, console_handler]
    )

    # 测试日志输出
    logger = logging.getLogger(__name__)

    coloredlogs.install(level='DEBUG', logger=logger,
                        fmt='%(asctime)s [%(levelname)s] %(message)s',
                        level_styles={
                            'debug': {'color': 'cyan'},
                            'info': {'color': 'green'},
                            'warning': {'color': 'yellow'},
                            'error': {'color': 'red'},
                            'critical': {'color': 'red', 'bold': True}
                        })

    logger.debug('调试信息（不可见）')
    logger.info('普通信息（出现在log文件和控制台）')
    logger.warning('警告信息（出现在log文件和控制台）')
    logger.error('错误信息（同时出现在正常log、错误log和控制台）')
    logger.critical('严重错误（同时出现在正常log、错误log和控制台）')

    logger.info('\n\n**********欢迎启动卓越智能体embedding平台**********\n\n')
except Exception as e:
    print("日志配置错误:")
    print(e)
    print(traceback.format_exc())



'''初始化后端api框架Fast-app'''

app = FastAPI()

'''跨域支持'''

# 设置允许的源列表
origins = ['*']  # * 允许所有
#     "http://localhost",
#     "http://localhost:8080",
#     "http://localhost:3000",
#     "https://your-frontend-domain.com",
# ]

# 添加 CORS 中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=origins,  # 允许的源列表
    allow_credentials=True,  # 允许携带凭证（如 cookies）
    allow_methods=["*"],  # 允许所有方法
    allow_headers=["*"],  # 允许所有头
)


'''api路由模块处理'''





'''统一总入参格式类定义'''

class publicarg(BaseModel):  # 公共参数，所有接口必传
    text: list = Field(frozen=True, description="文本内容[text, text]")
    apikey: str = Field(frozen=True, description="安全验证")


'''文本转向量智源开源模型加载'''

# 初始化模型（带检索指令）
model = FlagModel(mod_name, query_instruction_for_retrieval="为这个句子生成表示以用于检索相关文章：",
                normalize_embeddings=True,  # 开启归一化
                use_fp16=True  # 开户加速
                )



'''embedding文本转向量接口'''

@app.post("/embedding", tags=["embedding文本转向量接口"])
def t2v(mydata: publicarg):
    try:
        data_dict = mydata.model_dump()
        logger.warning(f'收到的请求={data_dict}')
        # 验证token、user
        if data_dict.get('apikey', '') != apikey:
            logger.warning(f'apikey验证失败')
            return {"msg": "apikey验证失败", "code": "403", "data": ""}
        # 处理cmd，返回对应的检索项数据
        if data_dict.get('text'):  # 有值才可以处理
            logger.warning(f'处理文本转向量')
            vdata = model.encode(data_dict.get('text')).tolist()
            print(type(vdata))
            return {"msg": "success", "code": "200", "data": vdata}
        else:  # cmd错误，无此检索项
            return {"msg": "text is null", "code": "404", "data": ""}

    except Exception as e:
        logger.error(f"embedding文本转向量接口: {e}")
        logger.error(traceback.format_exc())
        return {"msg": "error", "code": "501", "data": ""}



'''embedding短文本转向量接口，一般查询时用'''

@app.post("/embedding/queries", tags=["embedding短文本-转向量接口"])
def t2vq(mydata: publicarg):
    try:
        data_dict = mydata.model_dump()
        logger.warning('收到的请求')
        # 验证token、user
        if data_dict.get('apikey', '') != apikey:
            logger.warning(f'apikey验证失败')
            return {"msg": "apikey验证失败", "code": "403", "data": ""}
        # 处理cmd，返回对应的检索项数据
        if data_dict.get('text'):  # 有值才可以处理
            return {"msg": "success", "code": "200", "data": model.encode_queries(data_dict.get('text')).tolist()}
        else:  # cmd错误，无此检索项
            return {"msg": "text is null", "code": "404", "data": ""}

    except Exception as e:
        logger.error(f"embedding文本转向量接口: {e}")
        logger.error(traceback.format_exc())
        return {"msg": "error", "code": "501", "data": ""}



'''embedding文本相似度计算'''

@app.post("/embedding/similarity", tags=["embedding文本相似度计算接口"])
def similarity(mydata: publicarg):
    try:
        data_dict = mydata.model_dump()
        logger.warning('收到的请求')
        # 验证token、user
        if data_dict.get('apikey', '') != apikey:
            logger.warning(f'apikey验证失败')
            return {"msg": "apikey验证失败", "code": "403", "data": ""}
        # 处理cmd，返回对应的检索项数据
        if data_dict.get('text'):  # 有值才可以处理
            tv1 = model.encode(data_dict.get('text')[0])
            tv2 = model.encode(data_dict.get('text')[1])
            similarity = (tv1 @ tv2.T).tolist()
            logger.warning(f'相似度={similarity},{type(similarity)}')
            return {"msg": "success", "code": "200", "data": similarity}
        else:  # cmd错误，无此检索项
            return {"msg": "text is null", "code": "404", "data": ""}

    except Exception as e:
        logger.error(f"embedding文本转向量接口: {e}")
        logger.error(traceback.format_exc())
        return {"msg": "error", "code": "501", "data": ""}




# # 短查询自动添加指令
# queries = ["深度学习", "气候变化的影响"]
# query_embeddings = model.encode_queries(queries)  # 自动添加指令
#
# # 长文本无需指令（如文档库）
# passages = ["深度学习是机器学习的一个分支...", "气候变化导致全球气温上升..."]
# passage_embeddings = model.encode(passages)  # 不添加指令
#
# # 计算相似度
# similarity = query_embeddings @ passage_embeddings.T
# print(similarity)



# from FlagEmbedding import FlagAutoModel
#
# # pip install -U FlagEmbedding
#
# # D:\model\embeding\bge-small-zh-v1.5
# model = FlagAutoModel.from_finetuned('mod/bge-small-zh-v1.5', use_fp16=True)
#
# sentences_1 = "我喜欢太空, 想去那火星转一圈"
# sentences_2 = "该上班了, 我不喜欢上班"
# sentences_3 = "我喜欢星空, 想去太空旅游"
# embeddings_1 = model.encode(sentences_1)
#
# embeddings_2 = model.encode(sentences_2)
# embeddings_3 = model.encode(sentences_3)
# embeddings_4 = model.encode_queries(sentences_3)
#
# print('向量数据3=', embeddings_3)
#
# similarity = embeddings_1 @ embeddings_2.T
#
# similarity2 = embeddings_1 @ embeddings_3.T
#
# print(similarity)
#
# print(similarity2)




# listd = [sentences_1, sentences_2]
#
#
# # docs = [
# #     "呼叫中心功能包含通话录音，录音播放、录音下载功能。",
# #     "坐席管理可以增加坐席，修改密码，这是呼叫中心的坐席管理功能",
# #     "呼叫中心还有外呼记录、来电记录查询功能，可以看话单的通话时间等信息",
# # ]
#
# s12 = model.encode(listd)
#
# print(s12)

'''重排序'''

# from FlagEmbedding import FlagReranker
# reranker = FlagReranker('BAAI/bge-reranker-large', use_fp16=True) # Setting use_fp16 to True speeds up computation with a slight performance degradation
#
# score = reranker.compute_score(['query', 'passage'])
# print(score)
#
# scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']])
# print(scores)










